Data Analysis Assistant

Will LLMs be able to make my graphs prettier?
prompting
python
comparing models
Author

Grace

Published

April 6, 2025

Introduction

One of the important aspects of data analysis is making outputs that are easy to understand and aesthetically pleasing. In my work, I refer to this as making graphs “pretty”. Today I want to test whether LLMs have a concept of “prettiness”. I’m interested to see how they behave when I ask them to modify graph code to make it more aesthetically pleasing.

For reference, this is what the original graph looks like:

import pandas as pd
import plotly_express as px

anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

fig = px.choropleth(anx2010_df, locations="Code",
                    color="Anxiety_percent",
                    hover_name="Entity", 
                    color_continuous_scale=px.colors.sequential.Plasma)
fig.show()

Test 1: GPT 4 Turbo

from openai import OpenAI
import os
import pandas as pd

from IPython.display import Markdown, display

from dotenv import load_dotenv
# Load API key
_ = load_dotenv()

client = OpenAI()

model = "gpt-4-turbo-2024-04-09"
system_role = '''You are a consultant skilled at making visually appealing and easy to understand
graphs to present to clients.'''
initial_prompt = '''Improve this code to make the graph better. Reply only with Python code.

Here is the code to improve: 
import pandas as pd
import plotly_express as px

anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

fig = px.choropleth(anx2010_df, locations="Code",
                    color="Anxiety_percent",
                    hover_name="Entity", 
                    color_continuous_scale=px.colors.sequential.Plasma)
fig.show()
'''

messages = [
    {"role": "system", "content": system_role} ,
    {"role": "user", "content": initial_prompt},
]

completion = client.chat.completions.create(
  model=model,        # the model to use 
  messages=messages,  # a list of dictionaries with the prompt
  n=3,                # how many responses to generate
  temperature=0.7     # temperature parameter
)

response = display(Markdown(f'-----\n#### Response {0} \n {resp.message.content} \n'))

Response 0

import pandas as pd
import plotly_express as px

# Loading data
anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

# Creating a choropleth map with improved features
fig = px.choropleth(anx2010_df, locations="Code",
                   color="Anxiety_percent",
                   hover_name="Entity", 
                   title="Global Anxiety Rates in 2010",
                   labels={"Anxiety_percent": "Anxiety Rate (%)"},
                   color_continuous_scale=px.colors.sequential.Plasma,
                   projection="natural earth")

# Enhancing layout
fig.update_layout(
   geo=dict(
       showframe=False,
       showcoastlines=True,
       projection_type='equirectangular'
   ),
   coloraxis_colorbar=dict(
       title="Anxiety Rate %",
       thicknessmode="pixels", thickness=15,
       lenmode="pixels", len=300
   )
)

# Showing the plot
fig.show()
Testing the code
import pandas as pd
import plotly_express as px

# Loading data
anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

# Creating a choropleth map with improved features
fig = px.choropleth(anx2010_df, locations="Code",
                   color="Anxiety_percent",
                   hover_name="Entity", 
                   title="Global Anxiety Rates in 2010",
                   labels={"Anxiety_percent": "Anxiety Rate (%)"},
                   color_continuous_scale=px.colors.sequential.Plasma,
                   projection="natural earth")

# Enhancing layout
fig.update_layout(
   geo=dict(
       showframe=False,
       showcoastlines=True,
       projection_type='equirectangular'
   ),
   coloraxis_colorbar=dict(
       title="Anxiety Rate %",
       thicknessmode="pixels", thickness=15,
       lenmode="pixels", len=300
   )
)

# Showing the plot
fig.show()

Analysis

Here we can see that we have:

  1. A title for the overall graph
  2. A label for the anxiety rate legend
  3. No more black box around the outside of the map

Test 2: Claude 3.5 Sonnet

from dotenv import load_dotenv
from anthropic import Anthropic
import base64
import httpx
from IPython.display import display, Image, Markdown

# load api keys
_ = load_dotenv()
client = Anthropic()

message = client.messages.create(
    model='claude-3-5-sonnet-20241022',
    max_tokens=1000,
    temperature=0.0,
    system='''You are a consultant skilled at making visually appealing and easy to understand
graphs to present to clients.''',
    messages=[
        {"role": "user", "content": '''Improve this code to make the graph better. Reply only with Python code.

Here is the code to improve: 
import pandas as pd
import plotly_express as px

anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

fig = px.choropleth(anx2010_df, locations="Code",
                    color="Anxiety_percent",
                    hover_name="Entity", 
                    color_continuous_scale=px.colors.sequential.Plasma)
fig.show()
'''}
    ]
)

response = display(Markdown(f'-----\n#### Response {0} \n {message.content[0].text} \n'))

Response 0

import pandas as pd
import plotly_express as px

anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

fig = px.choropleth(anx2010_df, 
                   locations="Code",
                   color="Anxiety_percent",
                   hover_name="Entity",
                   hover_data={"Anxiety_percent": ":.2f"},
                   color_continuous_scale="RdYlBu_r",
                   title="Global Anxiety Prevalence (2010)",
                   labels={"Anxiety_percent": "Anxiety Prevalence (%)",
                          "Code": "Country Code",
                          "Entity": "Country"})

fig.update_layout(
   title_x=0.5,
   title_font_size=24,
   geo=dict(
       showframe=False,
       showcoastlines=True,
       projection_type='equirectangular'
   ),
   width=1200,
   height=700,
   margin=dict(t=50, l=0, r=0, b=0)
)

fig.update_coloraxes(
   colorbar_title_text="Prevalence (%)",
   colorbar_title_side="right",
   colorbar_thickness=15,
)

fig.show()

Testing the code

import pandas as pd
import plotly_express as px

anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

fig = px.choropleth(anx2010_df, 
                   locations="Code",
                   color="Anxiety_percent",
                   hover_name="Entity",
                   hover_data={"Anxiety_percent": ":.2f"},
                   color_continuous_scale="RdYlBu_r",
                   title="Global Anxiety Prevalence (2010)",
                   labels={"Anxiety_percent": "Anxiety Prevalence (%)",
                          "Code": "Country Code",
                          "Entity": "Country"})

fig.update_layout(
   title_x=0.5,
   title_font_size=24,
   geo=dict(
       showframe=False,
       showcoastlines=True,
       projection_type='equirectangular'
   ),
   width=1200,
   height=700,
   margin=dict(t=50, l=0, r=0, b=0)
)

fig.update_coloraxes(
   colorbar_title_text="Prevalence (%)",
   colorbar_title_side="right",
   colorbar_thickness=15,
)

fig.show()

Analysis

Here we see:

  1. A new color scheme
  2. Updated titles

Conclusion

Both models did pretty well with this task. A lot of a graph’s “prettiness” comes down to personal preference, but asking LLMs for various options could definitely help in the brainstorming process.